 /*
 * sample1.cpp -- Renders a solid green image.
 *
 * A filename can be given on the command line to write the results to file. 
 */

#include <iostream>
#include "random.h"
#include "sample1.h"
#include "kd.h"

using namespace optix; 

void Photon::initScene( InitialCameraData& camera_data ) {
	/* Primary RTAPI objects */ 
  Program ray_gen_program; 

  char path_to_ptx[512]; 

  width  = 512u;
  height = 512u; 
	// Photon caster parameters
	pl_width = 512u;
	pl_height = 512u;
  int use_glut = 1; 
	 
	_context->setRayTypeCount(4); 
	_context->setEntryPointCount(3); 

	 
	output_buffer = createOutputBuffer(RT_FORMAT_FLOAT4, width, height); 
  _context["result_buffer"]->setBuffer(output_buffer); 

	hit_buffer = _context->createBuffer(RT_BUFFER_INPUT_OUTPUT, RT_FORMAT_USER);
	hit_buffer->setElementSize(sizeof(ViewRayData));
	hit_buffer->setSize(width, height);
	_context["hit_buffer"]->setBuffer(hit_buffer);
	
	photon_buffer = _context->createBuffer(RT_BUFFER_INPUT_OUTPUT, RT_FORMAT_USER);
	photon_buffer->setElementSize(sizeof(PhotonRec));
	photon_buffer->setSize(pl_width*pl_height*maxPhotonsPerRay);
	_context["photon_buffer"]->setBuffer(photon_buffer);
	_context["maxPhotonsPerRay"]->setUint(maxPhotonsPerRay);
	_context["max_photon_depth"]->setUint(3);
	_context["flux_div_constant"]->setFloat(pl_width*pl_height / 4);

	kd_buffer = _context->createBuffer(RT_BUFFER_INPUT, RT_FORMAT_USER);
	kd_buffer->setElementSize(sizeof(PhotonRec));
	kd_buffer->setSize(KD::computeSize(pl_width, pl_height));
	_context["kd_buffer"]->setBuffer(kd_buffer);

	Buffer seed_buffer = _context->createBuffer(RT_BUFFER_INPUT, RT_FORMAT_UNSIGNED_INT2, width, height);
	srand(0);
	uint2* data = (uint2*) seed_buffer->map();
	for (int i=0; i<width; ++i) {
		for (int j=0; j<height; ++j) {
			data[i*height + j] = random2u(); //make_uint2(rand(), rand());
		}
	}
	seed_buffer->unmap();
	_context["seed_buffer"]->setBuffer(seed_buffer);

  sprintf( path_to_ptx, "%s/%s", sutilSamplesPtxDir(), "sample1_generated_plain.cu.ptx" );
	ray_gen_program = _context->createProgramFromPTXFile( path_to_ptx, "cast" ); 
	_context->setRayGenerationProgram(0, ray_gen_program); 

	sprintf( path_to_ptx, "%s/%s", sutilSamplesPtxDir(), "sample1_generated_photon.cu.ptx" );
	ray_gen_program = _context->createProgramFromPTXFile( path_to_ptx, "cast" );
	_context->setRayGenerationProgram(1, ray_gen_program);

	sprintf( path_to_ptx, "%s/%s", sutilSamplesPtxDir(), "sample1_generated_gather.cu.ptx" );
	ray_gen_program = _context->createProgramFromPTXFile( path_to_ptx, "cast" );
	_context->setRayGenerationProgram(2, ray_gen_program);

	/* Create geometry */ 
	Scene _scene(_context);

	/* Setup the camera */ 
 camera_data = InitialCameraData( make_float3( 278.0f, 273.0f, -800.0f ), // eye
                                   make_float3( 278.0f, 273.0f, 0.0f ),    // lookat
                                   make_float3( 0.0f, 1.0f,  0.0f ),       // up
                                   35.0f );      

	_context["eye"]->setFloat( make_float3( 0.0f, 0.0f, 0.0f ) );
	_context["U"]->setFloat( make_float3( 0.0f, 0.0f, 0.0f ) );
	_context["V"]->setFloat( make_float3( 0.0f, 0.0f, 0.0f ) );
	_context["W"]->setFloat( make_float3( 0.0f, 0.0f, 0.0f ) );	

  /* Run */
	_context->validate();
	_context->compile();

	_canRender = true;
}

// Glut main loop will call this with updated camera data
void  Photon::trace( const RayGenCameraData& camera_data ) {
	if (!_canRender) return;
	//_canRender = false;

	_context["eye"]->setFloat( camera_data.eye );
  _context["U"]->setFloat( camera_data.U );
  _context["V"]->setFloat( camera_data.V );
  _context["W"]->setFloat( camera_data.W );

	// Record hit locations and normals
	_context->launch(0, width, height); 
	// Trace photons
	_context->launch(1, pl_width, pl_height); 
	KD tree( (PhotonRec*) photon_buffer->map(), pl_width, pl_height, 
		make_float3(-160.0f, -1.5f, -1.0f), make_float3(560.0f, 550.0f, 560.0f) ); 
	photon_buffer->unmap();
	
	tree.store( (PhotonRec*)kd_buffer->map() );
	kd_buffer->unmap();
	
	_context->launch(2, width, height);

	//tree.gatherCPU( (ViewRayData *) hit_buffer->map(), (float *) output_buffer->map(), width, height, pl_width, pl_height );
	//hit_buffer->unmap();
	//output_buffer->unmap();
}

Buffer Photon::getOutputBuffer() {
	return output_buffer;
}

int main(int argc, char* argv[])
{
	Photon scene;
	GLUTDisplay::init( argc, argv );
	GLUTDisplay::run( "Photons", &scene, GLUTDisplay::CDNone );
    
  return( 0 );
}
 
